from datetime import datetime
import talib as ta
import pandas as pd
import matplotlib.pyplot as plt
import numpy as np
import time
import redis
from math import *
import traceback
from vnpy.trader.constant import Exchange, Interval
from vnpy.trader.object import HistoryRequest, LogData
from vnpy.trader.rqdata import rqdata_client
from vnpy.trader.database import database_manager
from vnpy.trader.ui import create_qapp, QtCore


def handle_close(evt):
    print("Closed Figure!")
    exit(0)


if __name__ == "__main__":
    # app = create_qapp()
    symbol = "IF2201"
    
    exchange = Exchange.CFFEX
    start = datetime(2021, 12, 1)
    end = datetime(2022, 3, 17)
    interval = Interval.MINUTE
    bars = database_manager.load_bar_data(
        symbol, Exchange.CFFEX, interval=Interval.MINUTE, start=start, end=end
    )

    if not bars or len(bars) == 0:
        df = pd.read_csv('IF2201_CFFEX_3056.csv', header=None)
        if df.empty or  len(df) == 0:
            req = HistoryRequest(
                symbol=symbol,
                exchange=exchange,
                interval=Interval.MINUTE,
                start=start,
                end=end,
            )
            try:
                bars = rqdata_client.query_history(req)

                if bars:
                    database_manager.save_bar_data(bars)
                    print(f"{symbol}-{interval}历史数据下载完成")
                else:
                    print(f"数据下载失败，无法获取{symbol}的历史数据")
            except Exception:
                msg = f"数据下载失败，触发异常：\n{traceback.format_exc()}"
                print(msg)

    n = 30
    history = bars[:n]
    left = len(bars) - n
    new_data = bars[n:]
    df = pd.DataFrame(
        [
            {
                "Close": s.close_price,
                "Low": s.low_price,
                "High": s.high_price,
                "Open": s.open_price,
                "Volume": s.volume,
            }
            for s in history
        ]
    )
    
    # df.to_csv('IF2201_CFFEX_3056.csv',index=False,sep=' ')
    # print(f"{df.index}:{df.values}")
    # # print(df.values)
    # plt.plot(df.index, df.values,'-r')
    # plt.show()#注意此函数需要调用
    # def update_bar():
    # for i in range(len(new_data)) :
    # def update_bar():
    #     if not new_data or len(new_data) <1:
    #         return
    #     # bar = new_data.pop(0)
    #     # print(new_data.pop(0).close_price)
    #     item = new_data.pop(0)

    #     df = df.append({'close_price':item.close_price},ignore_index= True)
    #     # df.add({len(df):new_data.pop(0).close_price})
    #     # print(df)
    #     # df.append(bar.close_price)
    #     plt.clf() #清空画布上的所有内容
    #     # plt.plot(df,'-r')
    #     plt.plot(df.index, df.values,'-r')
    #     plt.pause(0.01)
    fig = plt.figure()
    fig.canvas.mpl_connect("close_event", handle_close)
    df["up_band"], df["mid_band"], df["low_band"] = ta.BBANDS(
        df["Close"], timeperiod=20
    )
    last_tend = tend_value = 0
    offset_volume = 0
    pos = 0
    r = redis.Redis(host='localhost', port=6379, decode_responses=True)
    r.set(symbol ,  0)  # 设置 name 对应的值
    last = df.iloc[-1:n]
    print( f"%s%s" %(last['Close'],last['Close']))
    for i in range(len(new_data)):
        if not new_data or len(new_data) < 1:
            break
        # bar = new_data.pop(0)
        # print(new_data.pop(0).close_price)
        s = new_data.pop(0)

        df = df.append(
            {
                "Close": s.close_price,
                "Low": s.low_price,
                "High": s.high_price,
                "Open": s.open_price,
                "Volume": s.volume,
            },
            ignore_index=True,
        )

        # df.add({len(df):new_data.pop(0).close_price})
        # print(df)
        # df.append(bar.close_price)
        plt.clf()  # 清空画布上的所有内容
        # plt.plot(df,'-r')
        # plt.plot(df.index, df['Close'],'-r')
        # df["up_band"], df["mid_band"], df["low_band"] = ta.BBANDS(
        #     df["Close"], timeperiod=20
        # )
        # df["tend"] = ta.DEMA((df["Close"] - df["Open"]) * df["Volume"], timeperiod=30)
        # plt.plot(df.index, df['Volume'] +4700,'-g')
        # df[['Close','up_band','mid_band','low_band'] ].plot(figsize = (12,10))
        # plt.plot(df.index, df[['Close','up_band','mid_band','low_band'] ])
        df["Mean"] = ta.SMA(df["Close"], 30)
        
        sma = ta.SMA( abs(df['Close']-df['Open']) , 30)
        # sma = ta.SMA( abs(df['Close']-df['Open']) , self.slow_window)
        change = df['Close']-df['Open']
        vsma = ta.SMA(df['Volume'], 30)
        weight = df['Volume'] / (vsma +df['Volume'])
        # df['weight'] = df.apply(lambda row:  (row['Volume']+1)/(row['MeanVolume']+1)   ,axis =1)
        # close_sma_diff = df['Close'] - df["Mean"]
        boom = abs(change * weight) - weight * sma
        # print(boom[-1:] ) 
        if  float(boom[-1:]) >0 :
            print(df[-1:]) 

        df["MeanVolume"] = ta.SMA(df["Volume"], 30)
        df["Close2"] = df.apply(lambda row: float(row["Close"]) - 4900, axis=1)
        df["Flat"] = df.apply(
            lambda row: (float(row["Close"]) - float(row["Mean"])), axis=1
        )
        # df['Flat2'] = df.apply(lambda row:  (row['Volume'])- (row['MeanVolume']),axis =1)
        df["weight"] = df.apply(
            lambda row: (row["Volume"] + 1) / (row["MeanVolume"] + 1), axis=1
        )
        # df['weight'] = df.apply(lambda row:  (row['Volume']/(row['Volume'] + row['MeanVolume'])  ),axis =1)
        df["tend"] = ta.SMA(df["Flat"] , 10)
        df['tend'] = df.apply(lambda row:  row['tend'] *  row['weight']  ,axis =1)
        df['tend'] = df.apply(lambda row:  row['tend'] *  row['weight']  ,axis =1)
        df["one"] = 0
        # df['tend'] = ta.SMA(df['Flat'],5)
        # rowmax = df.max(axis=0)
        # n = df.size -1
        # print(df.size)
        # tend_value = df[rowmax]['tend']
        # print(s)
        # df['c'] = df.apply(lambda row: 1 if row['a'] < 0 and row['b'] > 0 else 0, axis=1)
        plt.plot(df.index, df[["Close2", "tend", "one"]])
        last = df.iloc [-1:]
        tend_value = float(last['tend'])
        if float(last_tend )*float(tend_value) <0 :
            key = symbol
            if float(tend_value) >0 :
                if pos > 0:
                    pass
                else:
                    offset_volume = 1 - pos
            else:
                if pos < 0:
                    pass
                else:
                    offset_volume = -1 - pos

            money = int(s.close_price *100* offset_volume)
            if r.exists(key):
                r.incrby(key , money)  # 设置 name 对应的值
            else:
                r.set(key ,  money)  # 设置 name 对应的值
            balance = r.get(key)
            
            pos = pos + offset_volume
            print(key + str(s.datetime), money, balance, pos,  "%s%s%s" %(last['weight'],last['tend'] ,last['Close']))
        # plt.plot(df.index,  1 ,label='linear line' )
        # r = redis.Redis(host='localhost', port=6379, decode_responses=True)
        # key = str(s.symbol)
        # print(f'tend_value = {tend_value}last_tend= {last_tend}')
        # if float(tend_value) * float(last_tend)< 0 :
        #     print(f'tend = {tend_value}price = {s.close_price}')
        # plt.subplot(df.index,df['Volume'])
        #     # plt.draw()#注意此函数需要调用
        #     # time.sleep(0.01)

        last_tend = tend_value
        plt.pause(0.01)

    # timer = QtCore.QTimer()
    # timer.timeout.connect(update_bar())
    # timer.start(1000)
    plt.show()
    # app.exec_()
